import ptan
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
from lib.efficient_kan import KAN
import gymnasium as gym

class AtariA2C(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(AtariA2C, self).__init__()

        # obs_action = (input_shape[2], input_shape[0], input_shape[1])
        print("obs_action: ", input_shape)
        self.conv = nn.Sequential(
            KAN([input_shape[0], 32, 64]),
            KAN([64, 64, 128]),
        )

        self.policy = nn.Sequential(
            KAN([128, 64, 512]),
            KAN([512, 64, n_actions])
        )

        self.value = nn.Sequential(
            KAN([128, 64, 512]),
            KAN([512, 64 , 1])
        )

    def forward(self, x):
        fx = x.float() / 255
        conv_out = self.conv(fx).view(fx.size()[0], -1)
        return self.policy(conv_out), self.value(conv_out)